import os
import time
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.common.exceptions import TimeoutException, NoSuchElementException

# 没有维护，抓取参数修改比较玄学，抓取也有bug,就是不存在的题目也会制作一个.md，存的是之前题目的内容
def setup_browser():
    options = Options()
    options.add_argument("--headless")  # 无头模式
    options.add_argument("--disable-gpu")  # 禁用GPU加速
    options.add_argument("--window-size=1920x1080")  # 设置窗口大小
    options.add_argument("--no-sandbox")
    browser = webdriver.Chrome(options=options)
    return browser

def calculate_directory(problem_id):
    """计算题目保存的目录"""
    number = int(problem_id[1:])  # 移除P或B前缀并转换为整数
    directory_number = number // 1000  # 除以1000，向下取整
    return str(directory_number)

def save_markdown(markdown_content, problem_id, base_path):
    """将Markdown内容保存为.md文件"""
    directory = os.path.join(base_path, problem_id[0], calculate_directory(problem_id))
    if not os.path.exists(directory):
        os.makedirs(directory)
    file_path = os.path.join(directory, f"{problem_id}.md")
    with open(file_path, 'w', encoding='utf-8') as file:
        file.write(markdown_content)

def main(base_path):
    browser = setup_browser()
    for prefix in ['B']:  # 遍历P和B开头的题目
        for i in range(3696, 4500, 1000):  # 设置题目范围
            for j in range(i, i + 1000):
                problem_id = f"{prefix}{j}"
                url = f"https://www.luogu.com.cn/problem/{problem_id}"
                try:
                    browser.get(url)
                    # 等待并点击复制Markdown的按钮
                    WebDriverWait(browser, 10).until(
                        EC.element_to_be_clickable((By.XPATH, "//*[@id='app']/div[2]/main/div/section[2]/section/div/div[1]/a[1]"))
                    )
                    copy_button = browser.find_element(By.XPATH, "//*[@id='app']/div[2]/main/div/section[2]/section/div/div[1]/a[1]")
                    copy_button.click()

                    # 使用JavaScript直接获取Markdown内容
                    markdown_content = browser.execute_script("return document.querySelector('textarea').value;")
                    save_markdown(markdown_content, problem_id, base_path)

                except (TimeoutException, NoSuchElementException):
                    print(f"Failed to load or find markdown copy button for {problem_id}")
                time.sleep(0.3)  # 避免请求过快
    browser.quit()

if __name__ == "__main__":
    # 将基础路径替换为你的存储地址
    main("D:/chen_education/project/SeleniumPractice/problem")
